#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import socket
import time
import threading
import logging
import sys
import os

# 尝试导入不同的键盘监听库
try:
    import keyboard  # 首选库
    KEYBOARD_LIB = 'keyboard'
except ImportError:
    try:
        import msvcrt  # Windows系统可用
        KEYBOARD_LIB = 'msvcrt'
    except ImportError:
        try:
            import tty
            import termios
            KEYBOARD_LIB = 'tty'
        except ImportError:
            print("错误: 无法导入任何键盘监听库。请安装keyboard库: pip install keyboard")
            sys.exit(1)

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class CommandGenerator:
    def __init__(self, server_host='127.0.0.1', server_port=8888, 
                 vx_step=0.1, vy_step=0.1, wz_step=0.1, 
                 max_vx=1.0, max_vy=1.0, max_wz=1.0):
        """
        初始化命令生成器
        :param server_host: 服务器地址
        :param server_port: 服务器端口
        :param vx_step: X方向速度步进值
        :param vy_step: Y方向速度步进值
        :param wz_step: 角速度步进值
        :param max_vx: X方向最大速度
        :param max_vy: Y方向最大速度
        :param max_wz: 最大角速度
        """
        self.server_host = server_host
        self.server_port = server_port
        self.vx_step = vx_step
        self.vy_step = vy_step
        self.wz_step = wz_step
        self.max_vx = max_vx
        self.max_vy = max_vy
        self.max_wz = max_wz
        
        # 当前速度值
        self.vx = 0.0
        self.vy = 0.0
        self.wz = 0.0
        
        # Socket连接
        self.socket = None
        self.connected = False
        self.running = False
        
        # 按键状态
        self.key_states = {
            'up': False,     # 前进
            'down': False,   # 后退
            'left': False,   # 左移
            'right': False,  # 右移
            'a': False,      # 左转
            'd': False       # 右转
        }
        
        # 按键映射
        self.key_mapping = {
            'up': {'key': 'up', 'action': self._increase_vx},
            'down': {'key': 'down', 'action': self._decrease_vx},
            'left': {'key': 'left', 'action': self._increase_vy},
            'right': {'key': 'right', 'action': self._decrease_vy},
            'a': {'key': 'a', 'action': self._increase_wz},
            'd': {'key': 'd', 'action': self._decrease_wz},
            'space': {'key': 'space', 'action': self._stop},
            'esc': {'key': 'esc', 'action': self._exit}
        }
    
    def connect(self):
        """
        连接到服务器
        """
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((self.server_host, self.server_port))
            self.connected = True
            logger.info(f"已连接到服务器 {self.server_host}:{self.server_port}")
            return True
        except Exception as e:
            logger.error(f"连接服务器失败: {e}")
            self.connected = False
            return False
    
    def _increase_vx(self):
        """增加X方向速度"""
        self.vx = min(self.vx + self.vx_step, self.max_vx)
        self._send_command()
    
    def _decrease_vx(self):
        """减少X方向速度"""
        self.vx = max(self.vx - self.vx_step, -self.max_vx)
        self._send_command()
    
    def _increase_vy(self):
        """增加Y方向速度"""
        self.vy = min(self.vy + self.vy_step, self.max_vy)
        self._send_command()
    
    def _decrease_vy(self):
        """减少Y方向速度"""
        self.vy = max(self.vy - self.vy_step, -self.max_vy)
        self._send_command()
    
    def _increase_wz(self):
        """增加角速度"""
        self.wz = min(self.wz + self.wz_step, self.max_wz)
        self._send_command()
    
    def _decrease_wz(self):
        """减少角速度"""
        self.wz = max(self.wz - self.wz_step, -self.max_wz)
        self._send_command()
    
    def _stop(self):
        """停止所有运动"""
        self.vx = 0.0
        self.vy = 0.0
        self.wz = 0.0
        self._send_command()
        logger.info("紧急停止")
    
    def _exit(self):
        """退出程序"""
        self._stop()
        self.running = False
        logger.info("退出程序")
    
    def _send_command(self):
        """发送控制命令"""
        if not self.connected:
            logger.warning("未连接到服务器，无法发送命令")
            return
        
        try:
            command = f"{self.vx:.2f},{self.vy:.2f},{self.wz:.2f}"
            self.socket.sendall(command.encode('utf-8'))
            logger.info(f"已发送命令: {command}")
        except Exception as e:
            logger.error(f"发送命令失败: {e}")
            self.connected = False
    
    def _handle_keyboard_input_keyboard_lib(self):
        """使用keyboard库处理键盘输入"""
        # 注册按键按下事件
        keyboard.on_press_key('up', lambda _: self._set_key_state('up', True))
        keyboard.on_press_key('down', lambda _: self._set_key_state('down', True))
        keyboard.on_press_key('left', lambda _: self._set_key_state('left', True))
        keyboard.on_press_key('right', lambda _: self._set_key_state('right', True))
        keyboard.on_press_key('a', lambda _: self._set_key_state('a', True))
        keyboard.on_press_key('d', lambda _: self._set_key_state('d', True))
        keyboard.on_press_key('space', lambda _: self._stop())
        keyboard.on_press_key('esc', lambda _: self._exit())
        
        # 注册按键释放事件
        keyboard.on_release_key('up', lambda _: self._set_key_state('up', False))
        keyboard.on_release_key('down', lambda _: self._set_key_state('down', False))
        keyboard.on_release_key('left', lambda _: self._set_key_state('left', False))
        keyboard.on_release_key('right', lambda _: self._set_key_state('right', False))
        keyboard.on_release_key('a', lambda _: self._set_key_state('a', False))
        keyboard.on_release_key('d', lambda _: self._set_key_state('d', False))
        
        # 保持线程运行
        while self.running:
            time.sleep(0.1)
    
    def _set_key_state(self, key, state):
        """设置按键状态"""
        if key in self.key_states:
            old_state = self.key_states[key]
            self.key_states[key] = state
            
            # 如果状态从False变为True，执行相应动作
            if not old_state and state:
                for mapping in self.key_mapping.values():
                    if mapping['key'] == key:
                        mapping['action']()
                        break
            
            # 如果状态从True变为False，停止相应方向的运动
            if old_state and not state:
                if key == 'up' or key == 'down':
                    self.vx = 0.0
                elif key == 'left' or key == 'right':
                    self.vy = 0.0
                elif key == 'a' or key == 'd':
                    self.wz = 0.0
                self._send_command()
    
    def _handle_keyboard_input_msvcrt(self):
        """使用msvcrt库处理键盘输入 (Windows)"""
        key_map = {
            b'H': 'up',      # 上箭头
            b'P': 'down',    # 下箭头
            b'K': 'left',    # 左箭头
            b'M': 'right',   # 右箭头
            b'a': 'a',       # a键
            b'd': 'd',       # d键
            b' ': 'space',   # 空格
            b'\x1b': 'esc'   # Esc
        }
        
        print("按键控制说明:")
        print("↑/↓: 前进/后退 (vx)")
        print("←/→: 左移/右移 (vy)")
        print("A/D: 左转/右转 (wz)")
        print("空格: 紧急停止")
        print("ESC: 退出程序")
        
        while self.running:
            if msvcrt.kbhit():
                key = msvcrt.getch()
                
                # 处理特殊键 (箭头键等)
                if key == b'\xe0':
                    key = msvcrt.getch()
                    if key in [b'H', b'P', b'K', b'M']:  # 箭头键
                        action_key = key_map.get(key)
                        if action_key:
                            for mapping in self.key_mapping.values():
                                if mapping['key'] == action_key:
                                    mapping['action']()
                                    break
                # 处理普通键
                elif key in [b'a', b'd', b' ', b'\x1b']:
                    action_key = key_map.get(key)
                    if action_key:
                        for mapping in self.key_mapping.values():
                            if mapping['key'] == action_key:
                                mapping['action']()
                                break
            
            time.sleep(0.1)
    
    def _handle_keyboard_input_tty(self):
        """使用tty库处理键盘输入 (Unix/Linux)"""
        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        
        print("按键控制说明:")
        print("w/s: 前进/后退 (vx)")
        print("a/d: 左移/右移 (vy)")
        print("q/e: 左转/右转 (wz)")
        print("空格: 紧急停止")
        print("ESC: 退出程序")
        
        try:
            tty.setraw(fd)
            while self.running:
                ch = sys.stdin.read(1)
                
                if ch == '\x1b':  # ESC键
                    self._exit()
                elif ch == ' ':  # 空格
                    self._stop()
                elif ch == 'w':  # 前进
                    self._increase_vx()
                elif ch == 's':  # 后退
                    self._decrease_vx()
                elif ch == 'a':  # 左移
                    self._increase_vy()
                elif ch == 'd':  # 右移
                    self._decrease_vy()
                elif ch == 'q':  # 左转
                    self._increase_wz()
                elif ch == 'e':  # 右转
                    self._decrease_wz()
                
                time.sleep(0.1)
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
    
    def start(self):
        """启动命令生成器"""
        if not self.connect():
            logger.error("无法启动命令生成器，连接服务器失败")
            return False
        
        self.running = True
        
        # 根据可用的库选择键盘输入处理方法
        if KEYBOARD_LIB == 'keyboard':
            print("使用keyboard库处理键盘输入")
            print("按键控制说明:")
            print("↑/↓: 前进/后退 (vx)")
            print("←/→: 左移/右移 (vy)")
            print("A/D: 左转/右转 (wz)")
            print("空格: 紧急停止")
            print("ESC: 退出程序")
            self._handle_keyboard_input_keyboard_lib()
        elif KEYBOARD_LIB == 'msvcrt':
            self._handle_keyboard_input_msvcrt()
        elif KEYBOARD_LIB == 'tty':
            self._handle_keyboard_input_tty()
        
        return True
    
    def stop(self):
        """停止命令生成器"""
        self.running = False
        
        # 关闭Socket连接
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
            self.socket = None
            self.connected = False
        
        logger.info("命令生成器已停止")

if __name__ == "__main__":
    import argparse
    
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='机器人控制命令生成器')
    parser.add_argument('--host', type=str, default='192.168.1.186', help='服务器地址')
    parser.add_argument('--port', type=int, default=8888, help='服务器端口')
    parser.add_argument('--vx-step', type=float, default=0.1, help='X方向速度步进值')
    parser.add_argument('--vy-step', type=float, default=0.1, help='Y方向速度步进值')
    parser.add_argument('--wz-step', type=float, default=0.1, help='角速度步进值')
    parser.add_argument('--max-vx', type=float, default=1.0, help='X方向最大速度')
    parser.add_argument('--max-vy', type=float, default=1.0, help='Y方向最大速度')
    parser.add_argument('--max-wz', type=float, default=1.0, help='最大角速度')
    args = parser.parse_args()
    
    # 创建并启动命令生成器
    generator = CommandGenerator(
        server_host=args.host,
        server_port=args.port,
        vx_step=args.vx_step,
        vy_step=args.vy_step,
        wz_step=args.wz_step,
        max_vx=args.max_vx,
        max_vy=args.max_vy,
        max_wz=args.max_wz
    )
    
    try:
        if generator.start():
            pass  # 主线程在start方法中阻塞
    except KeyboardInterrupt:
        logger.info("接收到退出信号")
    finally:
        generator.stop()
